package com.nbsaas.boot;

import javax.crypto.Cipher;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.util.Arrays;

public class AesNumHandle implements NumHandle {

    private static final String ALGO = "AES/GCM/NoPadding";
    private static final int GCM_TAG_BITS = 128;
    private static final int IV_LENGTH = 12; // 12 bytes for GCM
    private static final int KEY_LENGTH = 16; // AES-128

    private final byte[] aesKey;

    /**
     * 传入一个系统级密钥（可为任意字符串）
     */
    public AesNumHandle(String seedKey) {
        this.aesKey = deriveKey(seedKey);
    }

    private byte[] deriveKey(String seed) {
        try {
            MessageDigest digest = MessageDigest.getInstance("SHA-256");
            byte[] hash = digest.digest(seed.getBytes(StandardCharsets.UTF_8));
            return Arrays.copyOf(hash, KEY_LENGTH); // 取前16字节用于AES-128
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    private byte[] ivFromNum(int num) {
        try {
            MessageDigest digest = MessageDigest.getInstance("SHA-256");
            byte[] numBytes = ByteBuffer.allocate(4).putInt(num).array();
            byte[] hash = digest.digest(numBytes);
            return Arrays.copyOf(hash, IV_LENGTH);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public String encode(Integer num) {
        try {
            byte[] iv = ivFromNum(num);
            byte[] data = ByteBuffer.allocate(4).putInt(num).array();

            Cipher cipher = Cipher.getInstance(ALGO);
            GCMParameterSpec spec = new GCMParameterSpec(GCM_TAG_BITS, iv);
            cipher.init(Cipher.ENCRYPT_MODE, new SecretKeySpec(aesKey, "AES"), spec);

            byte[] cipherText = cipher.doFinal(data);
            return toHex(cipherText);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public Integer decode(String key) {
        try {
            byte[] cipherBytes = fromHex(key);

            // 因为 IV 可从明文数派生，但此处我们还没有明文；
            // 所以我们可逆方案需在加密输出中同时包含IV（可拼入前面）
            // 但你要求固定32位输出（16字节），所以需保持可推导IV。
            // 我们已用 num -> ivFromNum(num)，意味着IV取决于num，
            // 为了解密可逆，我们必须在密文中可还原出num的IV。
            // 因为IV基于num，不可从密文推导；所以必须加上IV。
            // 所以我们需要输出格式为：IV(12字节) + cipher(16字节) = 28字节 ≈ 56字符 hex。
            // 这时 decode 时才能取出IV再解密。

            // 因此需要调整 encode 实现，以下版本为同步调整后版本 ↓↓↓
            throw new UnsupportedOperationException("此版本 encode 未包含IV，无法解密。请看下一版。");
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    private String toHex(byte[] bytes) {
        StringBuilder sb = new StringBuilder(bytes.length * 2);
        for (byte b : bytes) {
            sb.append(String.format("%02x", b));
        }
        return sb.toString();
    }

    private byte[] fromHex(String s) {
        int len = s.length();
        byte[] data = new byte[len / 2];
        for (int i = 0; i < len; i += 2) {
            data[i / 2] = (byte) ((Character.digit(s.charAt(i), 16) << 4)
                    + Character.digit(s.charAt(i + 1), 16));
        }
        return data;
    }

    // 改进版：包含IV，保证可逆
    public String encodeWithIV(Integer num) {
        try {
            byte[] iv = ivFromNum(num);
            byte[] data = ByteBuffer.allocate(4).putInt(num).array();

            Cipher cipher = Cipher.getInstance(ALGO);
            GCMParameterSpec spec = new GCMParameterSpec(GCM_TAG_BITS, iv);
            cipher.init(Cipher.ENCRYPT_MODE, new SecretKeySpec(aesKey, "AES"), spec);

            byte[] cipherText = cipher.doFinal(data);

            // 拼接 iv + cipher
            byte[] combined = new byte[iv.length + cipherText.length];
            System.arraycopy(iv, 0, combined, 0, iv.length);
            System.arraycopy(cipherText, 0, combined, iv.length, cipherText.length);

            return toHex(Arrays.copyOf(combined, 16)); // 截断到16字节（如果必须32字符）
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }
}
